import numpy as np
from matplotlib import pyplot as plt
from matplotlib.ticker import (MultipleLocator, FormatStrFormatter, AutoMinorLocator)

def imshow_4D_rand(data_4D):
	print(data_4D.shape)

	nx = data_4D.shape[0]
	ny = data_4D.shape[1]
	nz = data_4D.shape[2]
	nt = data_4D.shape[3]

	index = np.unravel_index(data_4D.argmax(), data_4D.shape)
	print('max index is: ',index)
	nsx = index[0] # 32
	nsy = index[1] # 32
	nsz = index[2] # 32
	nst = index[3] # 16

    # grid for rand
	xg = np.arange(-640, 640, 20)
	yg = np.arange(-640, 640, 20)
	zg = np.arange(-1940, -660, 20)
	tg = np.arange(10, 330, 10)

	data_xz = np.squeeze(data_4D[:, nsy, :, nst])
	data_yz = np.squeeze(data_4D[nsx, :, :, nst])
	data_xy = np.squeeze(data_4D[:, :, nsz, nst])
	data_xt = np.squeeze(data_4D[:, nsy, nsz, :])
	print(data_xz.shape)
	print('max data_xz is %.2f'%(data_xz.max()))
	print('max data_yz is %.2f'%(data_yz.max()))
	print('max data_xy is %.2f'%(data_xy.max()))
	print('max data_xt is %.2f'%(data_xt.max()))

	plotx1 = 0.15
	plotx2 = 0.6
	ploty1 = 0.1
	ploty2 = 0.58
	plotw = 0.32
	ploth = 0.32
	
	font1 = {'size': 15}
	font2 = {'size': 13}
	font3 = {'size': 10}

	vmaxi = np.max(data_4D)
	if np.min(data_4D)<-0.4:
		vmini =-1 
	if np.min(data_4D)>=-0.4:
		vmini = 0 

	# print('max data 4D is: ',vmaxi)
	# print('min data 4D is: ',vmini)

	#plt.ion()
	fig = plt.figure(figsize = (7,6))
	ax1 = plt.axes([plotx1, ploty2, plotw, ploth])
	# im1 = plt.imshow(data_xz, cmap='viridis', vmax=1, vmin=-1, extent = [xg[0],xg[63],zg[0],zg[63]], origin = 'lower')
	im1 = plt.imshow(data_xz.transpose(), cmap='viridis', vmax=vmaxi, vmin=vmini, origin = 'lower')
	plt.xticks(np.arange(7,64,25), xg[7:64:25])
	plt.yticks(np.arange(7,64,25), zg[7:64:25])
	plt.tick_params(labelsize=12)
	ax1.set_aspect(1)
	plt.xlabel('x (m)', font2)
	plt.ylabel('z (m)', font2)
	plt.title('nsy=%d, nst=%d'%(nsy, nst), font1)
	# plt.text(900,-450,'data',size=20)
	# position=fig.add_axes([0.48, 0.6, 0.015, 0.1])
	# plt.colorbar(im1, cax=position, orientation='vertical')
	
	ax2 = plt.axes([plotx2, ploty2, plotw, ploth])
	# im2 = plt.imshow(data_yz, cmap='viridis', vmax=1, vmin=-1, extent = [yg[0],yg[63],zg[0],zg[63]], origin = 'lower')
	im2 = plt.imshow(data_yz.transpose(), cmap='viridis', vmax=vmaxi, vmin=vmini, origin = 'lower')
	plt.xticks(np.arange(7,64,25), yg[7:64:25])
	plt.yticks(np.arange(7,64,25), zg[7:64:25])
	plt.tick_params(labelsize=12)
	ax2.set_aspect(1)
	plt.xlabel('y (m)', font2)
	plt.ylabel('z (m)', font2)
	plt.title('nsx=%d, nst=%d'%(nsx, nst), font1)
	# position=fig.add_axes([0.48, 0.6, 0.015, 0.1])
	# plt.colorbar(im2, cax=position, orientation='vertical')
	 
	ax3 = plt.axes([plotx1, ploty1, plotw, ploth])
	# im3 = plt.imshow(data_xy, cmap='viridis', vmax=1, vmin=-1, extent = [xg[0],xg[63],yg[0],yg[63]], origin = 'lower')
	im3 = plt.imshow(data_xy.transpose(), cmap='viridis', vmax=vmaxi, vmin=vmini, origin = 'lower')
	plt.xticks(np.arange(7,64,25), xg[7:64:25])
	plt.yticks(np.arange(7,64,25), yg[7:64:25])
	plt.tick_params(labelsize=12)
	ax3.set_aspect(1)
	plt.xlabel('x (m)', font2)
	plt.ylabel('y (m)', font2)
	plt.title('nsz=%d, nst=%d'%(nsz, nst), font1)
	# position=fig.add_axes([0.48, 0.6, 0.015, 0.1])
	# plt.colorbar(im3, cax=position, orientation='vertical')
	
	ax4 = plt.axes([plotx2+0.02, ploty1+0.08, plotw/1.15, ploth])
	# im4 = plt.imshow(data_xt[0:32,:], cmap='viridis', vmax=1, vmin=-1, extent = [xg[0],xg[63],tg[0],tg[31]], origin = 'lower')
	im4 = plt.imshow(data_xt.transpose(), cmap='viridis', vmax=vmaxi, vmin=vmini, origin = 'lower')
	plt.xticks(np.arange(7,64,25), xg[7:64:25])
	plt.yticks(np.arange(4,32,10), tg[4:32:10]/1000)
	plt.tick_params(labelsize=12)
	ax4.set_aspect(1)
	plt.xlabel('x (m)', font2)
	plt.ylabel('t (s)', font2)
	plt.title('nsy=%d, nsz=%d'%(nsy, nsz), font1)
	position=fig.add_axes([0.62, 0.13, 0.28, 0.015])
	clb = plt.colorbar(im4, cax=position, orientation='horizontal')
	clb.ax.xaxis.set_major_locator(MultipleLocator(0.5*vmaxi))
	# clb.ax.yaxis.set_minor_locator(MultipleLocator(0.005))

	# plt.subplots_adjust(left=0.04, top= 0.90, right = 0.96, bottom = 0.12, wspace = 0.06, hspace = 0.44)
	
	plt.show()
	#plt.savefig('data_imshow.png')
	#plt.ioff()



def imshow_4D(data_4D):
	print(data_4D.shape)

	nx = data_4D.shape[0]
	ny = data_4D.shape[1]
	nz = data_4D.shape[2]
	nt = data_4D.shape[3]

	index = np.unravel_index(data_4D.argmax(), data_4D.shape)
	print('max index is: ',index)
	nsx = index[0] # 32
	nsy = index[1] # 32
	nsz = index[2] # 32
	nst = index[3] # 16

    # # grid for HS
	xg = np.arange(0, 64, 1)
	yg = np.arange(0, 64, 1)
	zg = np.arange(0, 64, 1)
	tg = np.arange(0, 32, 1)

	data_xz = np.squeeze(data_4D[:, nsy, :, nst])
	data_yz = np.squeeze(data_4D[nsx, :, :, nst])
	data_xy = np.squeeze(data_4D[:, :, nsz, nst])
	data_xt = np.squeeze(data_4D[:, nsy, nsz, :])
	print(data_xz.shape)
	print('max data_xz is %.2f'%(data_xz.max()))
	print('max data_yz is %.2f'%(data_yz.max()))
	print('max data_xy is %.2f'%(data_xy.max()))
	print('max data_xt is %.2f'%(data_xt.max()))

	plotx1 = 0.15
	plotx2 = 0.6
	ploty1 = 0.1
	ploty2 = 0.58
	plotw = 0.32
	ploth = 0.32
	
	font1 = {'size': 15}
	font2 = {'size': 13}
	font3 = {'size': 10}

	vmaxi = np.max(data_4D)
	if np.min(data_4D)<-0.4:
		vmini =-1 
	if np.min(data_4D)>=-0.4:
		vmini = 0 

	# print('max data 4D is: ',vmaxi)
	# print('min data 4D is: ',vmini)

	#plt.ion()
	fig = plt.figure(figsize = (7,6))
	ax1 = plt.axes([plotx1, ploty2, plotw, ploth])
	# im1 = plt.imshow(data_xz, cmap='viridis', vmax=1, vmin=-1, extent = [xg[0],xg[63],zg[0],zg[63]], origin = 'lower')
	im1 = plt.imshow(data_xz.transpose(), cmap='viridis', vmax=vmaxi, vmin=vmini, origin = 'lower')
	plt.xticks(np.arange(10,64,20), xg[10:64:20])
	plt.yticks(np.arange(10,64,20), zg[10:64:20])
	plt.tick_params(labelsize=12)
	ax1.set_aspect(1)
	plt.xlabel('x (m)', font2)
	plt.ylabel('z (m)', font2)
	plt.title('nsy=%d, nst=%d'%(nsy, nst), font1)
	# plt.text(900,-450,'data',size=20)
	# position=fig.add_axes([0.48, 0.6, 0.015, 0.1])
	# plt.colorbar(im1, cax=position, orientation='vertical')
	
	ax2 = plt.axes([plotx2, ploty2, plotw, ploth])
	# im2 = plt.imshow(data_yz, cmap='viridis', vmax=1, vmin=-1, extent = [yg[0],yg[63],zg[0],zg[63]], origin = 'lower')
	im2 = plt.imshow(data_yz.transpose(), cmap='viridis', vmax=vmaxi, vmin=vmini, origin = 'lower')
	plt.xticks(np.arange(10,64,20), yg[10:64:20])
	plt.yticks(np.arange(10,64,20), zg[10:64:20])
	plt.tick_params(labelsize=12)
	ax2.set_aspect(1)
	plt.xlabel('y (m)', font2)
	plt.ylabel('z (m)', font2)
	plt.title('nsx=%d, nst=%d'%(nsx, nst), font1)
	# position=fig.add_axes([0.48, 0.6, 0.015, 0.1])
	# plt.colorbar(im2, cax=position, orientation='vertical')
	 
	ax3 = plt.axes([plotx1, ploty1, plotw, ploth])
	# im3 = plt.imshow(data_xy, cmap='viridis', vmax=1, vmin=-1, extent = [xg[0],xg[63],yg[0],yg[63]], origin = 'lower')
	im3 = plt.imshow(data_xy.transpose(), cmap='viridis', vmax=vmaxi, vmin=vmini, origin = 'lower')
	plt.xticks(np.arange(10,64,20), xg[10:64:20])
	plt.yticks(np.arange(10,64,20), yg[10:64:20])
	plt.tick_params(labelsize=12)
	ax3.set_aspect(1)
	plt.xlabel('x (m)', font2)
	plt.ylabel('y (m)', font2)
	plt.title('nsz=%d, nst=%d'%(nsz, nst), font1)
	# position=fig.add_axes([0.48, 0.6, 0.015, 0.1])
	# plt.colorbar(im3, cax=position, orientation='vertical')
	
	ax4 = plt.axes([plotx2+0.02, ploty1+0.08, plotw/1.15, ploth])
	# im4 = plt.imshow(data_xt[0:32,:], cmap='viridis', vmax=1, vmin=-1, extent = [xg[0],xg[63],tg[0],tg[31]], origin = 'lower')
	im4 = plt.imshow(data_xt.transpose(), cmap='viridis', vmax=vmaxi, vmin=vmini, origin = 'lower')
	plt.xticks(np.arange(10,64,20), xg[10:64:20])
	plt.yticks(np.arange(10,32,10), tg[10:32:10])
	plt.tick_params(labelsize=12)
	ax4.set_aspect(1)
	plt.xlabel('x (m)', font2)
	plt.ylabel('t (s)', font2)
	plt.title('nsy=%d, nsz=%d'%(nsy, nsz), font1)
	position=fig.add_axes([0.62, 0.13, 0.28, 0.015])
	clb = plt.colorbar(im4, cax=position, orientation='horizontal')
	clb.ax.xaxis.set_major_locator(MultipleLocator(0.5*vmaxi))
	# clb.ax.yaxis.set_minor_locator(MultipleLocator(0.005))

	# plt.subplots_adjust(left=0.04, top= 0.90, right = 0.96, bottom = 0.12, wspace = 0.06, hspace = 0.44)
	
	plt.show()
	#plt.savefig('data_imshow.png')
	#plt.ioff()


